import numpy as np
import matplotlib.pyplot as plt


def fit_and_plot_average_flow_rate(speed_data, average_flow_rate_data, degree=3, show_plot=True, plot_filename=None):
    """
    对给定的转速和平均出水量数据进行多项式拟合，并可选地绘制和保存拟合曲线图。

    参数:
    - speed_data (list or array-like): 转速数据点
    - average_flow_rate_data (list or array-like): 平均出水量数据点
    - degree (int): 多项式的阶数，默认为3
    - show_plot (bool): 是否显示拟合结果图，默认为True
    - plot_filename (str or None): 如果提供，则将图表保存到指定文件名，默认不保存

    返回:
    - polynomial (np.poly1d object): 拟合得到的多项式对象
    - coefficients (array): 多项式系数，从最高次到最低次
    """
    # 将输入转换为numpy数组
    x = np.array(speed_data)
    y = np.array(average_flow_rate_data)

    # 拟合多项式
    coefficients = np.polyfit(x, y, degree)
    polynomial = np.poly1d(coefficients)

    if show_plot or plot_filename is not None:
        # 创建更精细的x轴数据用于绘图
        x_fit = np.linspace(x.min(), x.max(), 100)
        y_fit = polynomial(x_fit)

        plt.figure()
        plt.scatter(x, y, color='red', label='Average Data points')
        plt.plot(x_fit, y_fit, label=f'Fitted polynomial: {polynomial}')
        plt.title('Average Flow Rate vs Speed')
        plt.xlabel('Speed (rpm)')
        plt.ylabel('Average Flow Rate (L/min)')
        plt.legend()

        if plot_filename:
            plt.savefig(plot_filename)
        if show_plot:
            plt.show()
        else:
            plt.close()

    return polynomial, coefficients


# 使用直接提供的平均值数据进行拟合
speed_data = [50, 100, 150, 200, 250, 300, 350, 400, 450, 500]
average_flow_rate_data = [60.92, 118.02, 179.24, 246.69, 320.47, 385.27, 442.39, 483.20, 524.59, 567.05]

# 示例调用
if __name__ == "__main__":
    try:
        poly, coeffs = fit_and_plot_average_flow_rate(speed_data, average_flow_rate_data, degree=3, show_plot=True,
                                                      plot_filename="average_flow_rate_vs_speed.png")
        print("多项式系数（从最高次到最低次）:\n", coeffs)
    except Exception as e:
        print(f"Error encountered: {e}")